#pragma once
// HELPER
#include <algorithm>              // max

#include "../helper/helper.h"
#include "../helper/mpi_helper.h"
#include "multNMF.h"

namespace factor {
  template<typename MAT> class HALS : public MultNMF<MAT> {
  public:
    typedef MultNMF<MAT> __B;
    
  HALS(mpi::communicator& world_, po::variables_map& vm_) 
    : __B(world_, vm_) {}
    
    void NMF(MAT& A) {
      this->startTime = CLOCK;
      this->progress.clear();

      ARMA<MAT>::getSize(A, this->par.m, this->par.n);
      this->ProcessParameter();
    
      this->W = randu<mat>(this->par.m, this->par.k); // initialize
      this->H = randu<mat>(this->par.n, this->par.k);
    
      this->normA = norm(A, "fro");
      this->relE = 1;
      this->stop = false;
      FORE(iter, 0, this->par.maxiter) {
	mat AH = A*this->H, HTH = trans(this->H)*this->H;   // HALS for each column of W
	FORE(i, 0, this->par.k) {
	  this->W.col(i) = ALS(AH, HTH, this->W, i);
	}
	
	AH = trans(A)*this->W, HTH = trans(this->W)*this->W; // HALS for each column of H
	FORE(i, 0, this->par.k) {
	  this->H.col(i) = ALS(AH, HTH, this->H, i);
	}

	this->oldE = this->relE;                       // progress
	this->normE = norm(A-this->W*trans(this->H), "fro");
	this->relE = this->normE / this->normA;
	this->stop = this->StopCriteriaMet();          // stop condition
	double spW = ARMA<mat>::sparsity(this->W), spH = ARMA<mat>::sparsity(this->H);

	Iteration info(iter, this->normA, this->normE, this->relE, CLOCK - this->startTime, spW, spH, sqrt(this->normE*this->normE / this->par.m / this->par.n));
	this->progress.push_back(info);
	if (this->PrintCondition(iter)) info.Print();// output progress
	this->OutputProgress(info);

	if (this->stop) break;
      } // for iter

      if (!this->par.output.empty()) {     // output factors
	this->W.save(this->par.output+"_W");
	this->H.save(this->par.output+"_H");
      }
    } // nmf()

    void NMF(MAT& Arows, MAT& Acols) {
      this->startTime = CLOCK;
      this->progress.clear();

      ARMA<MAT>::getSize(Arows, this->mr, this->par.n);
      ARMA<MAT>::getSize(Acols, this->par.m, this->nc);
      this->ProcessParameter();
    
      this->ParInitialize();
      this->Combine(this->H, this->H_big);
      this->ComputeNormA(Arows);

      this->stop = false;
      this->relE = 1e-6;
      FORE(iter, 0, this->par.maxiter) {
	OptimizeBlock(Arows, Acols);

	ComputeNormE(Acols);
	this->oldE = this->relE;
	this->relE = this->normE / this->normA;
	if (MPI_IS_MASTER) this->stop = this->StopCriteriaMet();
	broadcast(this->world, this->stop, 0);

	if (MPI_IS_MASTER) {
	  double spW = ARMA<mat>::sparsity(this->W_big), spH = ARMA<mat>::sparsity(this->H_big);
	  Iteration info(iter, this->normA, this->normE, this->relE, CLOCK - this->startTime, spW, spH, sqrt(this->normE*this->normE / this->par.m / this->par.n));
	  this->progress.push_back(info);

	  if (this->PrintCondition(iter)) info.Print();
	  this->OutputProgress(info);
	}

	if (this->stop) break;
      }

      if (MPI_IS_MASTER && !this->par.output.empty()) {
	this->W_big.save(this->par.output+"_W");
	this->H_big.save(this->par.output+"_H");
      }
    }

    void OptimizeBlock(MAT& Arows, MAT& Acols) {
      mat AH = Arows*this->H_big, HTH = trans(this->H_big)*this->H_big;   // HALS for each column of W
      FORE(i, 0, this->par.k) {
	this->W.col(i) = ALS(AH, HTH, this->W, i);
      }
      this->Combine(this->W, this->W_big);

      AH = trans(Acols)*this->W_big, HTH = trans(this->W_big)*this->W_big; // HALS for each column of H
      FORE(i, 0, this->par.k) {
	this->H.col(i) = ALS(AH, HTH, this->H, i);
      }
      this->Combine(this->H, this->H_big);
    }
   
    vec ALS(const mat& AH, const mat& HTH, const mat& W, int i, double step = 1.0) {
      vec w = step * (AH.col(i) - W*HTH.col(i)) / HTH(i,i) + W.col(i);
      FORE(i, 0, (int) w.n_elem) {
	w(i) = max(w(i), this->par.epsilon);
      }
      return w;
    }
  }; // class HALS
 
} // namespace factor
